// Copyright Epic Games, Inc. All Rights Reserved.


#include "class06GameModeBase.h"
#include <Engine/StaticMeshActor.h>
#include <StaticMeshDescription.h>
#include <StaticMeshAttributes.h>
#include <EngineUtils.h>
#include <map>
#include <vector>
#include <fstream>
#include <iostream>
#include "lodepng.h"

void Aclass06GameModeBase::StartPlay()
{
    for (TActorIterator<AStaticMeshActor> itr(GetWorld()); itr; ++itr)
    {
        AStaticMeshActor* mesh = *itr;
        if (mesh->GetName() == "Table") AnalyseStaticMesh(mesh);
    }
    AGameModeBase::StartPlay();
}

void Aclass06GameModeBase::AnalyseStaticMesh(AStaticMeshActor* mesh)
{
    UStaticMeshComponent* component = mesh->GetStaticMeshComponent();
    UStaticMesh* meshData = component->GetStaticMesh();
    if (meshData)
    {
        int numLOD = meshData->GetNumLODs();
        int numVertices = meshData->GetNumVertices(0);
        int numTriangles = meshData->GetNumTriangles(0);
        UE_LOG(LogTemp, Warning, TEXT("%s: LOD = %d, Vertices = %d, Triangles = %d"),
            *meshData->GetName(), numLOD, numVertices, numTriangles);

        std::ofstream out("output.obj");
        std::vector<FVector3f> outPositions, outNormals;
        std::vector<FVector2f> outUVs;

        UStaticMeshDescription* desc = meshData->GetStaticMeshDescription(0);
        const TVertexAttributesRef<FVector3f>& positions = desc->GetVertexPositions();
        const TAttributesSet<FVertexInstanceID>& attributes = desc->VertexInstanceAttributes();
        TPolygonGroupAttributesConstRef<FName> materials = desc->GetPolygonGroupMaterialSlotNames();

        // Traverse all vertex attributes and triangles
        for (int i = 0; i < desc->GetPolygonGroupCount(); ++i)
        {
            TArray<FPolygonID> polygons;
            desc->GetPolygonGroupPolygons(i, polygons);
            UE_LOG(LogTemp, Warning, TEXT("%d: PolygonCount = %d, Material = %s"),
                i, polygons.Num(), *materials[i].ToString());

            for (FPolygonID pID : polygons)
            {
                TArray<FTriangleID> triangles;
                desc->GetPolygonTriangles(pID, triangles);
                for (FTriangleID tID : triangles)
                {
                    TArray<FVertexInstanceID> instanceIDs;
                    desc->GetTriangleVertexInstances(tID, instanceIDs);
                    for (int t = 0; t < 3; ++t)
                    {
                        FVertexInstanceID instanceID = instanceIDs[t];
                        FVertexID vID = desc->GetVertexInstanceVertex(instanceID);
                        FVector3f pos = positions[vID];
                        FVector3f norm = attributes.GetAttribute<FVector3f>(instanceID,
                            MeshAttribute::VertexInstance::Normal, 0);
                        FVector2f uv = attributes.GetAttribute<FVector2f>(instanceID,
                            MeshAttribute::VertexInstance::TextureCoordinate, 0);

                        outPositions.push_back(pos);
                        outNormals.push_back(norm);
                        outUVs.push_back(uv);
                    }
                }
            }
        }

        // Output materials
        FString mtlName;  // FIXME: handle only one group
        std::map<FString, FString> mtlFiles;
        for (int i = 0; i < desc->GetPolygonGroupCount(); ++i)
        {
            int mtlIndex = meshData->GetMaterialIndex(materials[i]);
            UMaterialInterface* mtl = meshData->GetMaterial(mtlIndex);
            mtlName = mtl->GetName();

#if true
            TArray<UTexture*> textures, textures2;
            mtl->GetTexturesInPropertyChain(
                EMaterialProperty::MP_BaseColor, textures, 
                NULL, NULL);
            mtl->GetTexturesInPropertyChain(
                EMaterialProperty::MP_Normal, textures2, 
                NULL, NULL);
            mtlFiles["map_Ka"] = AnalyseTextures(textures);
            mtlFiles["bump"] = AnalyseTextures(textures2);
#else
            TArray<UTexture*> textures;
            mtl->GetUsedTextures(textures, EMaterialQualityLevel::Num,
                                 false, ERHIFeatureLevel::Num, true);
            mtlFiles["map_Ka"] = AnalyseTextures(textures);
#endif
        }

        // Output to OBJ file
        out << "# Exported from Unreal: " << TCHAR_TO_UTF8(*mesh->GetName()) << std::endl;
        out << "mtllib output.mtl" << std::endl;
        out << "g " << TCHAR_TO_UTF8(*mtlName) << std::endl;
        out << "usemtl " << TCHAR_TO_UTF8(*mtlName) << std::endl;
        for (size_t i = 0; i < outPositions.size(); ++i)
        {
            const FVector3f& v = outPositions[i];
            out << "v " << v[0] << " " << v[1] << " " << v[2] << std::endl;
        }

        for (size_t i = 0; i < outNormals.size(); ++i)
        {
            const FVector3f& n = outNormals[i];
            out << "vn " << n[0] << " " << n[1] << " " << n[2] << std::endl;
        }

        for (size_t i = 0; i < outUVs.size(); ++i)
        {
            const FVector2f& uv = outUVs[i];
            out << "vt " << uv[0] << " " << uv[1] << std::endl;
        }

        for (size_t i = 0; i < outPositions.size(); i += 3)
            out << "f " << (i + 1) << "/" << (i + 1) << "/" << (i + 1) << " "
                << (i + 2) << "/" << (i + 2) << "/" << (i + 2) << " "
                << (i + 3) << "/" << (i + 3) << "/" << (i + 3) << " " << std::endl;
        UE_LOG(LogTemp, Warning, TEXT("Output to OBJ file: Vertices = %d"), (int)outPositions.size());

        // Output to MTL file
        std::ofstream out2("output.mtl");
        out2 << "# Exported from Unreal: " << TCHAR_TO_UTF8(*mesh->GetName()) << std::endl;
        out2 << "newmtl " << TCHAR_TO_UTF8(*mtlName) << std::endl;
        out2 << "Ka 0.2 0.2 0.2" << std::endl;
        out2 << "Kd 0.6 0.6 0.6" << std::endl;
        out2 << "Ks 0.9 0.9 0.9" << std::endl;
        for (std::map<FString, FString>::iterator itr = mtlFiles.begin();
             itr != mtlFiles.end(); ++itr)
        {
            out2 << TCHAR_TO_UTF8(*itr->first) << " "
                 << TCHAR_TO_UTF8(*itr->second) << std::endl;
        }
    }
}

FString Aclass06GameModeBase::AnalyseTextures(TArray<UTexture*> textures)
{
    FString resultFileName;
    for (UTexture* tex : textures)
    {
        UTexture2D* tex2D = dynamic_cast<UTexture2D*>(tex);
        if (!tex2D) continue; // we only handle 2D textures

        std::vector<unsigned char> outImageData;
        int w = tex2D->GetSizeX(), h = tex2D->GetSizeY();
        UE_LOG(LogTemp, Warning, TEXT("Texture %s: %d x %d"),
               *tex2D->GetFName().ToString(), w, h);

        // Calling LockReadOnly may get a nullptr in return if the texture
        // is not in a specific format as below:
        // 1. CompressionSettings set to VectorDisplacementmap
        // 2. MipGenSettings to NoMipmaps
        // 3. SRGB to false
        TextureCompressionSettings prevCompression = tex2D->CompressionSettings;
        TextureMipGenSettings prevMipSettings = tex2D->MipGenSettings;
        uint8 prevSRGB = tex2D->SRGB;
        tex2D->CompressionSettings = TextureCompressionSettings::TC_VectorDisplacementmap;
        tex2D->MipGenSettings = TextureMipGenSettings::TMGS_NoMipmaps;
        tex2D->SRGB = false;
        tex2D->UpdateResource();

        const FColor* imageData = static_cast<const FColor*>(
            tex2D->PlatformData->Mips[0].BulkData.LockReadOnly());
        if (imageData != NULL)
        {
            for (int y = 0; y < h; ++y)
                for (int x = 0; x < w; ++x)
                {
                    FColor color = imageData[y * w + x];
                    outImageData.push_back(color.R);
                    outImageData.push_back(color.G);
                    outImageData.push_back(color.B);
                    outImageData.push_back(color.A);
                }
        }
        tex2D->PlatformData->Mips[0].BulkData.Unlock();

        std::string fileName = TCHAR_TO_UTF8(*tex2D->GetFName().ToString());
        unsigned int result = lodepng::encode(fileName + ".png", outImageData, w, h);
        if (result > 0)
        {
            UE_LOG(LogTemp, Warning, TEXT("lodepng: %s"),
                   UTF8_TO_TCHAR(lodepng_error_text(result)));
        }
        else
        {
            UE_LOG(LogTemp, Warning, TEXT("lodepng: %s, %d"),
                   *tex2D->GetFName().ToString(), (int)outImageData.size());
        }
        resultFileName = UTF8_TO_TCHAR((fileName + ".png").c_str());  // FIXME: handle only one file

        tex2D->CompressionSettings = prevCompression;
        tex2D->MipGenSettings = prevMipSettings;
        tex2D->SRGB = prevSRGB;
        tex2D->UpdateResource();
    }
    return resultFileName;
}
